import os
import pickle
import json
import cv2
import numpy as np
from scipy.spatial.transform import Rotation as Rt
import open3d as o3d


def read_pose_from_json(pose_path, num=0):
    raw_data = json.load(open(pose_path, "r"))
    if "dataList" in raw_data:
        raw_pose = raw_data["dataList"][num]
    else:
        raw_pose = raw_data["objects"][num]
    
    translation, rotation, scale = raw_pose["center"], raw_pose["rotation"], raw_pose["dimensions"]
    translation = np.float32([translation["x"], translation["y"], translation["z"]])
    rotation = np.float32([rotation["x"], rotation["y"], rotation["z"]])
    rotation = Rt.from_euler('XYZ', rotation).as_matrix()
    scale = np.float32([scale["length"], scale["width"], scale["height"]])

    pose = {
        "translation": translation.reshape(3, 1),
        "rotation": rotation,
        "scale": scale.reshape(3, 1),
    }
    return pose
